import numpy as np
from PIL import Image, ImageDraw

from ship_envs.core.eitity import Entity
from ship_envs.core.model import Model


class ObsViewer(object):
    def __init__(self, width, height):
        self.width = width
        self.height = height
        self.xfact = self.width # 物理模型宽高默认为1*1，此处需对模型进行放大，放大因子为需显示的宽高
        self.yfact = self.height
        self.models = []

    def add_model(self, model):
        if model and isinstance(model, Model):
            self.models.append(model)

    def add_entity(self, entity):
        if entity and isinstance(entity, Entity):
            self.models.append(entity.get_collision_model())

    def _new_img(self):
        img = Image.new("RGB", (self.width, self.height), (255, 255, 255))
        drawer = ImageDraw.Draw(img)
        for model in self.models:
            drawer.polygon(model.get_scaled_pts(self.xfact, self.yfact), tuple((model.color * 255).astype(np.int32)))
        return img

    def get_array(self, clear=True):
        img = self._new_img()
        data = np.array(img.getdata()).reshape(img.width, img.height, 3)
        data = np.transpose(data, (2, 0, 1))
        img.close()
        if clear:
            self.clear_models()
        return data

    def clear_models(self):
        self.models = []

    def get_img(self, clear=True):
        img = self._new_img()
        if clear:
            self.clear_models()
        return img
